# /**
# * Copyright (c) 2022 Beijing Jiaotong University
# * PhotLab is licensed under [Open Source License].
# * You can use this software according to the terms and conditions of the [Open Source License].
# *
# * You may obtain a copy of [Open Source License] at: [https://open.source.license/]
# *
# * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# *
# * See the [Open Source License] for more details.
# */
# /**
# * Author: Shihao Xie
# * Created: Sep. 3, 2022
# * Supported by: National Key Research and Development Program of China
# */
import os

import numpy as np
from phot_engine import ZWEngine, ModelOfComputation
from phot_engine.kernel.manager.task_manager import TaskManager

from phot_server.manager.graph_manager import GraphManager
from phot import logger
from phot_engine import TaskPriority
from phot_engine import TaskStatus

from phot_server.repository.project_config_repository import ProjectConfigRepository
from phot_server.rpc.engine_data_type.ttypes import PhotTaskStatus, PhotOutput, PhotComplexDouble, PhotSignal, \
    PhotArray, \
    PhotArrayType, PhotAnalyzerOutput, PhotPlotType
from phot_server.utils.toml_parser import save_toml, parse_toml
from phot_server.utils.type_convert import object_to_dict, PhotGraph_to_GraphEntity


class EngineManager:
    def __init__(self):
        self.__graph_manager = GraphManager()
        self.computing_engine = ZWEngine()
        self.task_manager = TaskManager.get_instance()
        self.__project_repository = ProjectConfigRepository()
        self.__task_graph_map = {}

    async def add_task(self, graph, render_graph, task_id, priority=None):
        entity = PhotGraph_to_GraphEntity(graph, render_graph)
        # graph_dict = object_to_dict(graph)
        # entity = graph_dict2GraphEntity(graph_dict)
        # entity.print_graph_entity_meta_data(show_detail=True)
        if priority is None:
            await entity.submit_simulation_task(task_id=task_id, priority=TaskPriority.LOW,moc=ModelOfComputation.SEQUENCE, fast_simulation=False)
        else:
            await entity.submit_simulation_task(task_id=task_id, priority=priority,moc=ModelOfComputation.SEQUENCE, fast_simulation=False)
        # if execution_time is not None:
        #     task.set_end_executing_time(execution_time)
        # task.show_task_info()
        # task.show_computational_graph()
        # self.computing_engine._task_manager.add_task(task_id, task)
        self.__task_graph_map[task_id] = graph

    # add a task to engine manager
    # @param graph_name: the graph to be calculated
    # @param task_id: the task id
    # @param priority: the task priority
    # @param execution_time: the task execution time
    # @return PhotResponse: add task response
    # exception PhotError: ...
    def create_task_from_existing_graph(self, graph_id, task_id, priority):
        graph = self.__graph_manager.get_graph(graph_id)
        render_graph = self.__graph_manager.get_render_graph(graph_id)
        self.__task_graph_map[task_id] = graph
        return self.add_task(graph, render_graph, task_id, priority)

    # engine calculate a task
    # @param task_id: the task id to be calculated
    # @return PhotResponse: calculate task result response
    # exception PhotError: ...
    def calculate_task(self, task_id):
        self.computing_engine.calculate_task(task_id)
        graph = self.__task_graph_map[task_id]
        if self.get_task_status(task_id) == PhotTaskStatus.TASK_ACCOMPLISHED:
            project_config = self.__project_repository.get_project_config_by_name(graph.name)
            from phot_server.server_config.global_config import TEMP_ROOT
            path = TEMP_ROOT
            if project_config is not None:
                for file in project_config.files:
                    if file.name == "Outputs":
                        path = file.path if file is not None else TEMP_ROOT
            for component in graph.components:
                if component.isAnalyzer:
                    self.save_analyzer_result_as(task_id, component.componentId,
                                                 os.path.join(path, str(component.componentId) + ".toml"))
                else:
                    self.save_component_output_as(task_id, component.componentId,
                                                  os.path.join(path, str(component.componentId) + ".toml"))

    # cancel a task to engine
    # @param graph_name: the graph to be canceled
    # @param task_id: the task id
    # @return PhotResponse: cancel task response
    # exception PhotError: ...
    def cancel_task(self, task_id):
        self.computing_engine.cancel_task(task_id)

    # # cancel a task to engine
    # # @param graph_name: the graph to be paused
    # # @param task_id: the task id
    # # @return PhotResponse: cancel task response
    # # exception PhotError: ...
    def pause_task(self, task_id):
        self.computing_engine.pause_task(task_id)

    # recover a task
    # @param graph_name: the graph to be recovered
    # @param task_id: the task id
    # @return PhotResponse: recover task response
    # exception PhotError: ...
    def recover_task(self, task_id):
        self.computing_engine.recover_task(task_id)

    # delete a task
    # @param graph_name: the graph to be deleted
    # @param task_id: the task id
    # @return PhotResponse: recover task response
    # exception PhotError: ...
    def delete_task(self, task_id):
        self.computing_engine.delete_task(task_id)
        self.__task_graph_map.pop(task_id)

    def restart_task(self, task_id):
        # self.computing_engine.cancel_task(task_id)
        self.computing_engine.delete_task(task_id)
        self.create_task_from_existing_graph(self.__task_graph_map[task_id].graphId, task_id, None)
        self.computing_engine.calculate_task(task_id)

    def get_all_component_output(self, task_id):
        result_edge_list = self.computing_engine.fetch_all_simulation_result(task_id=task_id)
        # component_signal_map = {}
        # for edge in result_edge_list:
        #     source = edge.source_port
        #     source_name = "{}_port:{}".format(source.vertex_id, source.port_id)
        #     signal = edge.get_signal()
        #     component_signal_map[source_name] = str(signal)
        # np.set_printoptions(threshold=10)
        return result_edge_list

    def signal_to_PhotArray(self, signal):
        if signal.dtype == np.dtype(np.complex128):
            return PhotArray(complexDoubleArray=PhotComplexDouble(real=signal.reshape(-1).real.tolist(),
                                                                  imag=signal.reshape(-1).imag.tolist())), \
                   PhotArrayType.complexDoubleArray
        elif signal.dtype == np.dtype(np.float64):
            return PhotArray(doubleArray=signal.reshape(-1).tolist()), PhotArrayType.doubleArray
        elif signal.dtype == np.dtype(np.int32):
            return PhotArray(intArray=signal.reshape(-1).tolist()), PhotArrayType.intArray
        elif signal.dtype == np.dtype(np.int64):
            return PhotArray(longArray=signal.reshape(-1).tolist()), PhotArrayType.longArray

    def get_component_output(self, task_id, component_id):
        outputList = []
        for signal_info in self.computing_engine.get_task(task_id).get_signal_by_vertex_id(str(component_id)):
            if isinstance(signal_info["signal"], list):
                polarization_num = len(signal_info["signal"])
            else:
                polarization_num = 1
            output = PhotOutput(componentId=int(signal_info["source_vertex_id"]),
                                portId=int(signal_info["source_port_id"]),
                                polarization=str(polarization_num),
                                signalType="",
                                signals=[])
            for polarization in range(polarization_num):
                if isinstance(signal_info["signal"], list):
                    signal = signal_info["signal"][polarization]
                    signal_shape = signal.shape
                else:
                    continue

                # print("Signal_info : type {}  -  shape {}".format(signal_type, signal_shape))
                phot_signal = PhotSignal(shape=str(signal_shape),
                                         orientation="X" if polarization == 0 else "Y")
                phot_signal.signal, phot_signal.type = self.signal_to_PhotArray(signal)
                output.signals.append(phot_signal)
            outputList.append(output)
        return outputList

    def get_analyzer_outputs(self, task_id, component_id):
        analyzer_signal_list = self.computing_engine.get_task(task_id).get_signal_by_vertex_id(str(component_id))
        graph = self.__task_graph_map.get(task_id, None)
        analyzer_plotType_map = {
            "constellation_diagram": PhotPlotType.SCATTER,
            "eye_diagram": PhotPlotType.LINECHART,
            "bits_error_count": PhotPlotType.TABLE
        }
        for signal_info in analyzer_signal_list:
            source_component_id = signal_info["source_vertex_id"]
            signal = signal_info["signal"]
            taget_analyzer = None
            for component in graph.components:
                if source_component_id == str(component.componentId):
                    taget_analyzer = component
                    break
            if taget_analyzer is None:
                return None
            analyzer_outputs = []
            if type(signal) == int:
                analyzer_output = PhotAnalyzerOutput(graphId=graph.graphId, analyzerId=component_id,
                                                     plotType=analyzer_plotType_map[taget_analyzer.calculationMethod],
                                                     title=taget_analyzer.name)
                analyzer_output.shape = 1
                analyzer_output.x = PhotArray(intArray=[signal])
                analyzer_outputs.append(analyzer_output)
            elif type(signal) == list:
                if type(signal[0]) == list:
                    for i in range(len(signal)):
                        analyzer_output = PhotAnalyzerOutput(graphId=graph.graphId, analyzerId=component_id,
                                                             plotType=analyzer_plotType_map[
                                                                 taget_analyzer.calculationMethod],
                                                             title=taget_analyzer.name + str(i + 1))
                        analyzer_output.shape = str(signal[i][0].shape)
                        analyzer_output.x = self.signal_to_PhotArray(signal[i][0])[0]
                        analyzer_output.y = self.signal_to_PhotArray(signal[i][1])[0]
                        analyzer_outputs.append(analyzer_output)
                else:
                    analyzer_output = PhotAnalyzerOutput(graphId=graph.graphId, analyzerId=component_id,
                                                         plotType=analyzer_plotType_map[
                                                             taget_analyzer.calculationMethod],
                                                         title=taget_analyzer.name)
                    analyzer_output.shape = str(signal[0].shape)
                    analyzer_output.x = self.signal_to_PhotArray(signal[0])[0]
                    analyzer_output.y = self.signal_to_PhotArray(signal[1])[0]
                    analyzer_outputs.append(analyzer_output)
            elif type(signal) == np.ndarray:
                analyzer_output = PhotAnalyzerOutput(graphId=graph.graphId, analyzerId=component_id,
                                                     plotType=analyzer_plotType_map[taget_analyzer.calculationMethod],
                                                     title=taget_analyzer.name)
                analyzer_output.shape = str(signal)
                analyzer_output.matrix = signal.tolist()
                analyzer_outputs.append(analyzer_output)
            return analyzer_outputs

    # save a component result in a task
    # @param task_id: the task id
    # @param component_id: the component id to be got output
    # @param path: file path to output
    # @return PhotResponse: recover task response
    # exception PhotError: ...

    def save_component_output_as(self, task_id, component_id, file_path):
        outputs = self.get_component_output(task_id, component_id)
        output_list = []
        for output in outputs:
            output_dict = object_to_dict(output)
            output_list.append(output_dict)
        save_toml(file_path, {"PhotOutput": output_list})

    # get a task result
    # @param task_id: the task id
    # @param component_id: the component id to be got output
    # @param directory_path: directory path to output
    # @return PhotResponse: recover task response
    # exception PhotError: ...
    def save_analyzer_result_as(self, task_id, component_id, file_path):
        outputs = self.get_analyzer_outputs(task_id, component_id)
        output_dict = {"PhotAnalyzerOutput": []}
        for output in outputs:
            output_dict["PhotAnalyzerOutput"].append(object_to_dict(output))
        save_toml(file_path, output_dict)

    # get a task status
    # @param task_id: the task id
    # @return PhotResponse: recover task response
    # exception PhotError: ...
    def get_task_status(self, task_id):
        status = self.computing_engine.get_task_status(task_id)
        if status == TaskStatus.TASK_CREATING:
            return PhotTaskStatus.TASK_CREATING
        elif status == TaskStatus.TASK_RUNNING:
            return PhotTaskStatus.TASK_RUNNING
        elif status == TaskStatus.TASK_STOPPED:
            return PhotTaskStatus.TASK_STOPPED
        elif status == TaskStatus.TASK_NO_RESPONDING:
            return PhotTaskStatus.TASK_NO_RESPONDING
        elif status == TaskStatus.TASK_ACCOMPLISHED:
            return PhotTaskStatus.TASK_ACCOMPLISHED
        elif status == TaskStatus.TASK_DIED:
            return PhotTaskStatus.TASK_DIED
        return None

    # delete a task
    # @param graph_name: the graph to be deleted
    # @param task_id: the task id
    # @return PhotResponse: recover task response
    # exception PhotError: ...
    # def get_task_exception_information(self, task_id):
    #     all_result = self.computing_engine.fetch_all_simulation_result(task_id=task_id)
    #     status = self.get_task_status(task_id)
    #     if exception_queue.empty():
    #         exception = "None"
    #         exception_list = []
    #     else:
    #         exception = "None"
    #         exception_list = []
    #         while not exception_queue.empty():
    #             exception = repr(exception_queue.get())
    #             exception_list.append(exception)
    #     if len(exception_list) > 1:
    #         logger.error(exception_list)
    #         return exception_list
    #     elif status == PhotTaskStatus.TASK_ACCOMPLISHED:
    #         return ["The task {0} is finished without exception!".format(task_id)]
    #     else:
    #         message = "The task {0} exception is {1}.".format(task_id, exception)
    #         return [message]

    def fetch_all_simulation_result(self,task_id):
        all_result = self.computing_engine.fetch_all_simulation_result(task_id=task_id)
        print("Fetch all simulation result: ")
        for result in all_result:
            print(result)
        return all_result

    def get_history_output(self, task_id, component_id):
        from phot_server.server_config.global_config import TEMP_ROOT
        graph = self.__task_graph_map.get(task_id, None)
        if graph is None:
            return None
        project_config = self.__project_repository.get_project_config_by_name(graph.name)
        path = TEMP_ROOT
        if project_config is not None:
            for file in project_config.files:
                if file.name == "Outputs":
                    path = file.path if file.path is not None else TEMP_ROOT
        file_name = str(component_id) + ".toml"
        if file_name in os.listdir(path):
            output_dict = parse_toml(os.path.join(path, file_name))
            if "PhotOutput" in output_dict.keys():
                data = output_dict["PhotOutput"]
                output = PhotOutput(componentId=data["componentId"], portId=data["componentId"],
                                    polarization=data["componentId"],
                                    signalType=data["componentId"])
            elif "PhotAnalyzerOutput" in output_dict.keys():
                pass
            return None
